import os
import subprocess
import tkinter as tk
from tkinter import scrolledtext, messagebox, filedialog, ttk
import threading

# 安装必要的库
def install_requirements():
    try:
        import huggingface_hub
    except ImportError:
        print("huggingface_hub 库未安装，正在安装...")
        os.system("pip install huggingface_hub -i https://pypi.tuna.tsinghua.edu.cn/simple/")
        from huggingface_hub import list_repo_files

install_requirements()

try:
    from huggingface_hub import list_repo_files
except ImportError:
    print("huggingface_hub 库未安装，请手动安装。")
    exit(1)

# 显示模型/数据集清单
def list_files(repo_id, repo_type):
    try:
        return list_repo_files(repo_id=repo_id, repo_type=repo_type)
    except Exception as e:
        print(f"Error listing files for {repo_id}: {str(e)}")
        return []

# 模型/数据集下载主程序
def download_model_or_dataset(model, token, include, exclude, dataset, save_dir, use_hf_transfer, use_mirror):
    token_option = f"--token {token}" if token else ""
    include_option = f"--include {include}" if include else ""
    exclude_option = f"--exclude {exclude}" if exclude else ""

    if model:
        model_name = model.split("/")
        save_dir_option = (
            f"--local-dir {os.path.join(save_dir, 'models--' + model_name[0] + '--' + model_name[1])}"
            if save_dir
            else ""
        )

        download_command = (
            f"huggingface-cli download {token_option} {include_option} {exclude_option} "
            f"--local-dir-use-symlinks False --resume-download {model} {save_dir_option}"
        )
    elif dataset:
        dataset_name = dataset.split("/")
        save_dir_option = (
            f"--local-dir {os.path.join(save_dir, 'datasets--' + dataset_name[0] + '--' + dataset_name[1])}"
            if save_dir
            else ""
        )

        download_command = (
            f"huggingface-cli download {token_option} {include_option} {exclude_option} "
            f"--local-dir-use-symlinks False --resume-download --repo-type dataset {dataset} {save_dir_option}"
        )

    if use_hf_transfer:
        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
        print("HF transfer enabled")

    if use_mirror:
        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
        print("Mirror enabled")

    subprocess.run(download_command, shell=True)

# 更新选择的模型/数据集清单
def update_file_list(repo_id, repo_type, file_list_var):
    files = list_files(repo_id, repo_type)
    file_list_var.set(files)

# 开始下载
def start_download(model, token, include, exclude, dataset, save_dir, use_hf_transfer, use_mirror, status_label):
    def download_thread():
        status_label.config(text="下载开始...")
        download_model_or_dataset(model, token, include, exclude, dataset, save_dir, use_hf_transfer, use_mirror)
        status_label.config(text="下载完成！")

    thread = threading.Thread(target=download_thread)
    thread.start()

# 创建主窗口
root = tk.Tk()
root.title("Huggingface模型/数据集GUI下载工具")
root.geometry("800x600")  # 设置初始窗口大小
root.resizable(True, True)  # 允许调整窗口大小

# 创建 Notebook (Tab 控件)
notebook = ttk.Notebook(root)
notebook.pack(fill=tk.BOTH, expand=True)

# 模型设置 Tab
model_tab = ttk.Frame(notebook)
notebook.add(model_tab, text="模型设置")

# 数据集设置 Tab
dataset_tab = ttk.Frame(notebook)
notebook.add(dataset_tab, text="数据集设置")

# 模型设置 Tab 内容
model_frame = tk.Frame(model_tab)
model_frame.pack(pady=10, padx=10, fill=tk.BOTH, expand=True)

model_label = tk.Label(model_frame, text="模型名称:")
model_label.grid(row=0, column=0, sticky=tk.W)
model_entry = tk.Entry(model_frame, width=60)
model_entry.grid(row=0, column=1, padx=5, pady=5)

update_button = tk.Button(model_frame, text="更新文件列表(需能直连Huggingface)", command=lambda: update_file_list(model_entry.get(), "model", file_list_var))
update_button.grid(row=1, column=0, columnspan=2, pady=5)

file_list_var = tk.StringVar(value=[])
file_list = tk.Listbox(model_frame, listvariable=file_list_var, selectmode=tk.MULTIPLE, height=10)
file_list.grid(row=2, column=0, columnspan=2, pady=5)

# 数据集设置 Tab 内容
dataset_frame = tk.Frame(dataset_tab)
dataset_frame.pack(pady=10, padx=10, fill=tk.BOTH, expand=True)

dataset_label = tk.Label(dataset_frame, text="数据集名称:")
dataset_label.grid(row=0, column=0, sticky=tk.W)
dataset_entry = tk.Entry(dataset_frame, width=60)
dataset_entry.grid(row=0, column=1, padx=5, pady=5)

# 公共设置
common_frame = tk.Frame(root)
common_frame.pack(pady=10, padx=10, fill=tk.BOTH, expand=True)

token_label = tk.Label(common_frame, text="token:")
token_label.grid(row=0, column=0, sticky=tk.W)
token_entry = tk.Entry(common_frame, width=60)
token_entry.grid(row=0, column=1, padx=5, pady=5)

include_label = tk.Label(common_frame, text="包含的下载文件 (可选):")
include_label.grid(row=1, column=0, sticky=tk.W)
include_entry = tk.Entry(common_frame, width=60)
include_entry.grid(row=1, column=1, padx=5, pady=5)

exclude_label = tk.Label(common_frame, text="排除的下载文件 (可选):")
exclude_label.grid(row=2, column=0, sticky=tk.W)
exclude_entry = tk.Entry(common_frame, width=60)
exclude_entry.grid(row=2, column=1, padx=5, pady=5)

save_dir_label = tk.Label(common_frame, text="保存目录(必选):")
save_dir_label.grid(row=3, column=0, sticky=tk.W)
save_dir_entry = tk.Entry(common_frame, width=60)
save_dir_entry.grid(row=3, column=1, padx=5, pady=5)
save_dir_button = tk.Button(common_frame, text="选择目录", command=lambda: save_dir_entry.insert(0, filedialog.askdirectory()))
save_dir_button.grid(row=3, column=2, padx=5, pady=5)

use_hf_transfer_var = tk.BooleanVar(value=False)
use_hf_transfer_check = tk.Checkbutton(common_frame, text="使用hf_transfer加速下载", variable=use_hf_transfer_var)
use_hf_transfer_check.grid(row=4, column=0, columnspan=3, pady=5, sticky=tk.W)

use_mirror_var = tk.BooleanVar(value=True)
use_mirror_check = tk.Checkbutton(common_frame, text="从HuggingFace镜像站下载", variable=use_mirror_var)
use_mirror_check.grid(row=5, column=0, columnspan=3, pady=5, sticky=tk.W)

# 下载按钮和状态标签
download_button = tk.Button(root, text="下载", command=lambda: start_download(
    model_entry.get(),
    token_entry.get(),
    include_entry.get(),
    exclude_entry.get(),
    dataset_entry.get(),
    save_dir_entry.get(),
    use_hf_transfer_var.get(),
    use_mirror_var.get(),
    status_label
))
download_button.pack(pady=10, padx=10, fill=tk.X, expand=True)

status_label = tk.Label(root, text="")
status_label.pack(pady=10, padx=10, fill=tk.X, expand=True)

# 运行主循环
root.mainloop()
